
#include "select_v2_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling
{
    const uint32_t BLOCK_SIZE = 32;
    const uint32_t BUFFER_NUM = 2;
    static ge::graphStatus TilingFunc(gert::TilingContext *context)
    {

        SelectV2TilingData tiling;

        uint32_t x0Size = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
        uint32_t x1Size = context->GetInputShape(1)->GetStorageShape().GetShapeSize();
        uint32_t x2Size = context->GetInputShape(2)->GetStorageShape().GetShapeSize();
        uint32_t ySize = context->GetOutputShape(0)->GetStorageShape().GetShapeSize();

        uint32_t tilingkey = 0;
        if (ySize != x0Size || ySize != x1Size || ySize != x2Size)
        {
            tilingkey = 5;

            int32_t y_ndarray[20], x0_ndarray[20], x1_ndarray[20], x2_ndarray[20];
            int32_t y_dimensional, x0_dimensional, x1_dimensional, x2_dimensional;
            auto shape_y = context->GetOutputShape(0)->GetOriginShape();
            auto shape_x0 = context->GetInputTensor(0)->GetOriginShape();
            auto shape_x1 = context->GetInputTensor(1)->GetOriginShape();
            auto shape_x2 = context->GetInputTensor(2)->GetOriginShape();

            y_dimensional = shape_y.GetDimNum();
            x0_dimensional = shape_x0.GetDimNum();
            x1_dimensional = shape_x1.GetDimNum();
            x2_dimensional = shape_x2.GetDimNum();

            int32_t max_y_dimensional;
            max_y_dimensional = y_dimensional;
            if (x1_dimensional > max_y_dimensional)
                max_y_dimensional = x1_dimensional;
            if (x2_dimensional > max_y_dimensional)
                max_y_dimensional = x2_dimensional;
            if (x0_dimensional > max_y_dimensional)
                max_y_dimensional = x0_dimensional;

            for (int i = 0; i < max_y_dimensional; i++)
            {
                if (i < y_dimensional)
                    y_ndarray[y_dimensional - i - 1] = shape_y.GetDim(i);
                else
                    y_ndarray[i] = 1;

                if (i < x0_dimensional)
                    x0_ndarray[x0_dimensional - i - 1] = shape_x0.GetDim(i);
                else
                    x0_ndarray[i] = 1;

                if (i < x1_dimensional)
                    x1_ndarray[x1_dimensional - i - 1] = shape_x1.GetDim(i);
                else
                    x1_ndarray[i] = 1;

                if (i < x2_dimensional)
                    x2_ndarray[x2_dimensional - i - 1] = shape_x2.GetDim(i);
                else
                    x2_ndarray[i] = 1;
            }
            tiling.set_y_dimensional(max_y_dimensional);
            tiling.set_y_ndarray(y_ndarray);
            tiling.set_x0_ndarray(x0_ndarray);
            tiling.set_x1_ndarray(x1_ndarray);
            tiling.set_x2_ndarray(x2_ndarray);

            int32_t y_sumndarray[20], x0_sumndarray[20], x1_sumndarray[20], x2_sumndarray[20];
            y_sumndarray[0] = 1;
            x0_sumndarray[0] = 1;
            x1_sumndarray[0] = 1;
            x2_sumndarray[0] = 1;
            for (int i = 1; i <= max_y_dimensional; i++)
            {
                y_sumndarray[i] = y_sumndarray[i - 1] * y_ndarray[i - 1];
                x0_sumndarray[i] = x0_sumndarray[i - 1] * x0_ndarray[i - 1];
                x1_sumndarray[i] = x1_sumndarray[i - 1] * x1_ndarray[i - 1];
                x2_sumndarray[i] = x2_sumndarray[i - 1] * x2_ndarray[i - 1];
            }
            tiling.set_y_sumndarray(y_sumndarray);

            tiling.set_x0_sumndarray(x0_sumndarray);
            tiling.set_x1_sumndarray(x1_sumndarray);
            tiling.set_x2_sumndarray(x2_sumndarray);
        }
        else
        {
            uint64_t ubSize;
            auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
            ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);

            uint32_t x1Num = context->GetInputShape(1)->GetStorageShape().GetShapeSize();
            uint32_t x1Bytes = GetSizeByDataType(context->GetInputDesc(1)->GetDataType());
            uint32_t x1Length = x1Num * x1Bytes;

            uint32_t ubDataNumber = 5;
            ge::DataType dataType = context->GetInputTensor(1)->GetDataType();
            switch (dataType)
            {
            case ge::DT_FLOAT:
                ubDataNumber = 5;
                tilingkey = 1;
                break;
            case ge::DT_FLOAT16:
                ubDataNumber = 5;
                tilingkey = 2;
                break;
            case ge::DT_INT32:
                ubDataNumber = 6;
                tilingkey = 3;
                break;
            case ge::DT_INT8:
                ubDataNumber = 9;
                tilingkey = 4;
                break;
            default:
                break;
            }

            uint32_t tileBlockNum = (ubSize / BLOCK_SIZE / BUFFER_NUM) / ubDataNumber;
            uint32_t tileDataNum = tileBlockNum * BLOCK_SIZE / x1Bytes;
            uint32_t inputLengthAlgin32 = ((x1Length + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE;
            uint32_t everyCoreInputBlockNum = inputLengthAlgin32 / BLOCK_SIZE;
            uint32_t coreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / x1Bytes;
            uint32_t tileNum = everyCoreInputBlockNum / tileBlockNum;
            uint32_t finalTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? tileNum : tileNum + 1;
            uint32_t tailDataNum = coreDataNum - tileDataNum * tileNum;
            tailDataNum = tailDataNum == 0 ? tileDataNum : tailDataNum;

            tiling.set_coreDataNum(coreDataNum);   
            tiling.set_tileDataNum(tileDataNum);   
            tiling.set_finalTileNum(finalTileNum); 
            tiling.set_tailDataNum(tailDataNum);  
        }
        context->SetBlockDim(1);
        context->SetTilingKey(tilingkey);
        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

        size_t *currentWorkspace = context->GetWorkspaceSizes(1);
        currentWorkspace[0] = 0;

        return ge::GRAPH_SUCCESS;
    }
}

namespace ge
{
    static ge::graphStatus InferShape(gert::InferShapeContext *context)
    {
        const gert::Shape *x1_shape = context->GetInputShape(0);
        gert::Shape *y_shape = context->GetOutputShape(0);
        *y_shape = *x1_shape;
        return GRAPH_SUCCESS;
    }
}

namespace ops
{
    class SelectV2 : public OpDef
    {
    public:
        explicit SelectV2(const char *name) : OpDef(name)
        {
            this->Input("condition")
                .ParamType(REQUIRED)
                .DataType({ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
            this->Input("x1")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_INT8})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
            this->Input("x2")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_INT8})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
            this->Output("y")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_INT8})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

            this->SetInferShape(ge::InferShape);

            this->AICore()
                .SetTiling(optiling::TilingFunc);
            this->AICore().AddConfig("ascend310b").AddConfig("ascend910b");
        }
    };

    OP_ADD(SelectV2);
}
